package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"math/rand"
	"net/http"
	"time"
	"unicode/utf8"
)

var UpGrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

var TarMap = make(map[string]*websocket.Conn)

var connHost *websocket.Conn

func HandWsOne(c *gin.Context) error {
	ws, err := UpGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		fmt.Println(err)
		return err
	}
	input := struct {
		TargetId string `form:"targetId"`
	}{}
	if err = c.ShouldBindQuery(&input); err != nil {
		return err
	}
	if _, ok := TarMap[input.TargetId]; !ok {
		fmt.Println("目标连接不存在")
		return err
	}
	connHost = ws
	for {
		var message []byte
		_, message, err = ws.ReadMessage()
		if err != nil {
			fmt.Println(err)
			return err
		}
		fmt.Println("被控端,say:", string(message))
		if utf8.Valid(message) {
			fmt.Println("是文本消息")
			err = TarMap[input.TargetId].WriteMessage(websocket.TextMessage, message)
		} else {
			fmt.Println("是二进制消息")
			err = TarMap[input.TargetId].WriteMessage(websocket.BinaryMessage, message)
		}
		if err != nil {
			fmt.Println(err)
			return err
		}
	}
}

func HandWsTar(c *gin.Context) error {
	ws, err := UpGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		fmt.Println(err)
		return err
	}
	id := RandString(1)
	TarMap[id] = ws
	err = ws.WriteMessage(websocket.TextMessage, []byte(id))
	if err != nil {
		fmt.Println(err)
		return err
	}
	for {
		if TarMap[id] != nil {
			var msg []byte
			_, msg, err = TarMap[id].ReadMessage()
			if err != nil {
				fmt.Println(err)
				return err
			}
			fmt.Println("tar say:", msg)
			if utf8.Valid(msg) {
				fmt.Println("是文本消息")
				err = connHost.WriteMessage(websocket.TextMessage, msg)
			} else {
				fmt.Println("是二进制消息")
				err = connHost.WriteMessage(websocket.BinaryMessage, msg)
			}
			if err != nil {
				return err
			}
		}
	}
}

func RandString(len int) string {
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	bytes := make([]byte, len)
	for i := 0; i < len; i++ {
		b := r.Intn(26) + 65
		bytes[i] = byte(b)
	}
	return string(bytes)
}
